//
// Created by tacom on 22-9-1.
//
#include "opencv2/cudaimgproc.hpp"
#include "CNFD.cuh"
#include "SegCudaCommon.h"
#include "set"

void colorLabelsCpu(const cv::Mat &labels, cv::Mat &colors) {
    for (int r = 0; r < labels.rows; ++r) {
        const int * labels_row = labels.ptr<int>(r);
        cv::Vec3b * colors_row = colors.ptr<cv::Vec3b>(r);
        for (int c = 0; c < labels.cols; ++c) {
            colors_row[c] = cv::Vec3b(labels_row[c] * 131 % 255, labels_row[c] * 241 % 255, labels_row[c] * 251 % 255);
        }
    }
}

__global__ void colorLabelGpu(cv::cuda::PtrStepSz<int> labels, cv::cuda::PtrStepSz<uchar3> colors){
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    int j = blockDim.y * blockIdx.y + threadIdx.y;

    if (i < labels.cols && j < labels.rows){
        uchar3 tmp;
        int tmpValue = labels(j, i);
        tmp.x = tmpValue * 131 % 255;
        tmp.y = tmpValue * 241 % 255;
        tmp.z = tmpValue * 251 % 255;
        colors(j, i) = tmp;
    }
}


unsigned int findMatMaxValue(cv::cuda::GpuMat &gpuMat){
    // 这个是统计联通域数量的函数, cpu是按照0-max来标记的，gpu则按照线程号，所以这个命名有误
    // 不像cpu从0开始编号，这里是根线程id来编号的, 并且没有很好的在gpu中执行的map库，所以下载到cpu中统计
    cv::Mat cpuMat;
    gpuMat.download(cpuMat);

    std::set<int> historyMat;
    for(int i = 0; i < cpuMat.rows; ++i){
        const int *cpuMatRows = cpuMat.ptr<int>(i);
        for(int j = 0; j < cpuMat.cols; ++j){
            const int value = cpuMatRows[j];
            if(historyMat.find(value) == historyMat.end()){
                historyMat.insert(value);
            }
        }
    }
    return historyMat.size() - 1;
}


void GetConnected(cv::cuda::GpuMat &inputMat, cv::cuda::GpuMat &outputMat) {
    cv::cuda::GpuMat binaryMat;
    cv::cuda::cvtColor(inputMat, binaryMat, cv::COLOR_BGR2GRAY);
    cv::cuda::connectedComponents(binaryMat, outputMat);
}


void CNFDCudaWrapper(cv::cuda::GpuMat &inputMat, cv::cuda::GpuMat &outputMat,
                     int h, int w, int stride, std::vector<std::vector<int>> &matrixShape,
                     std::vector<unsigned int> &CNFDList) {
    // 输出结果
    outputMat.create(inputMat.size(), CV_8UC3);

    try{
        for (int i = 0; i < matrixShape.size(); ++i){
            for (int j = 0; j < matrixShape[0].size(); ++j){
                if(matrixShape[i][j] == 0) return ;

                // cal roi
                int hStart = i * h + i * stride;
                int wStart = j * w + j * stride;
                cv::Range height(hStart, hStart + h);
                cv::Range width(wStart, wStart + w);
                cv::cuda::GpuMat eachInRoi(inputMat, height, width);
                cv::cuda::GpuMat eachToRoi;

                // cal connected
                GetConnected(eachInRoi, eachToRoi);

                // output
                cv::cuda::GpuMat eachOutRoi(outputMat, height, width);

                dim3 threadDim(32, 32);
                dim3 blockDim(SegCudaCommon::calBlockSize(eachToRoi.cols, threadDim.x),
                              SegCudaCommon::calBlockSize(eachToRoi.rows, threadDim.y));

                // output img
                colorLabelGpu<<<blockDim, threadDim>>>(eachToRoi, eachOutRoi);

                // output CNFD
                CNFDList.push_back(findMatMaxValue(eachToRoi));
            }
        }
    }catch (const std::exception &e){
        std::cout << e.what() << std::endl;
    }

}